# main.py

import time

from langchain.memory import ConversationBufferMemory
from langchain_mistralai import ChatMistralAI
from Chat.model_handler import Text2SQLGenerator, SQL2TextGenerator
from Chat.prompt_templates import SQL_PROMPT, RESULT_PROMPT
from Tools.database import UniversalDBConnector
from typing import Dict, Any
import re
import json
import os


def Text2SQL(llm, memory):
    """自然语言转换SQL语句"""
    try:
        return Text2SQLGenerator(llm, SQL_PROMPT, memory)
    except Exception as e:
        print("Text2SQL初始化错误:", e)
        return None


def SQL2Text(llm):
    """SQL结果转自然语言"""
    try:
        return SQL2TextGenerator(llm, RESULT_PROMPT)
    except Exception as e:
        print("SQL2Text初始化错误:", e)
        return None


def parse_json_response(response):
    """增强版JSON解析"""
    try:
        json_match = re.search(r"```json(.*?)```", response, re.DOTALL)
        if not json_match:
            raise ValueError("未找到JSON代码块")

        json_str = json_match.group(1).strip()
        data = json.loads(json_str)

        # 如果iscode存在的话，检查其值是否为true
        if "iscode" in data and data["iscode"] == "true":
            data["iscode"] = True

        return data
    except Exception as e:
        raise ValueError(f"JSON解析失败: {str(e)}")


def call_ai_for_final_answer(interpreter, question, result, code):
    """获取最终回答"""
    return interpreter.generate_text({
        "question": question,
        "result": str(result),  # 确保结果为字符串
        "sqlcode": code,
    })


def call_ai_for_sql(generator, db, question):
    """获取SQL查询"""
    return generator.generate_sql({
        "dialect": db.dialect,
        "table_info": db.get_table_info(),
        "question": question
    })


def main():
    # 初始化组件
    llm = ChatMistralAI(
        api_key=os.getenv("MISTRAL_API_KEY"),
        model="codestral-2501",
        temperature=0.3
    )

    # 使用新版Memory初始化
    memory = ConversationBufferMemory(
        memory_key="history",
        input_key="question",
        output_key="answer",
        return_messages=True
    )

    db = UniversalDBConnector("sqlite:///student.db")
    sql_generator = Text2SQL(llm, memory)
    result_interpreter = SQL2Text(llm)

    if not sql_generator or not result_interpreter:
        print("系统初始化失败")
        return

    while True:
        try:
            question = input("\n用户问题（输入exit退出）: ").strip()
            if question.lower() == "exit":
                break
            if not question:
                continue

            # 生成SQL响应
            sql_response = call_ai_for_sql(sql_generator, db, question)
            # print(f"\n[原始响应]\n{sql_response}")

            sql_response = re.sub(r'<think>.*?</think>', '', sql_response, flags=re.DOTALL).strip()

            # 解析响应
            sql_data = parse_json_response(sql_response)
            sql_code = sql_data["sqlcode"].strip()
            is_code = sql_data["iscode"]

            if is_code and sql_code:
                print(f"\n[生成SQL]\n{sql_code}")

                # 执行安全查询
                result = db.execute_safe_query(sql_code)

                if sql_code.strip().upper().startswith("DELETE") or sql_code.strip().upper().startswith(
                        "UPDATE") or sql_code.strip().upper().startswith("INSERT"):
                    # 对于DELETE操作，检查受影响的行数
                    affected_rows = result.rowcount
                    if affected_rows > 0:
                        # print(f"\n[提示] 已删除 {affected_rows} 条记录")
                        time.sleep(2)

                        # 生成最终解释
                        final_response = call_ai_for_final_answer(
                            result_interpreter,
                            question,
                            f" 已处理 {affected_rows} 条记录",
                            sql_code
                        )
                        final_response = re.sub(r'<think>.*?</think>', '', final_response, flags=re.DOTALL).strip()
                        final_data = parse_json_response(final_response)
                        print(f"\n[最终回答]\n{final_data['message']}")

                        memory.save_context(
                            {"question": question},
                            {"answer": f"已删除 {affected_rows} 条记录"}
                        )
                    else:
                        print("\n[提示] 未删除任何记录")
                        memory.save_context(
                            {"question": question},
                            {"answer": "未删除任何记录"}
                        )
                else:
                    rows = result.fetchall()
                    if rows:
                        # print(rows)
                        # print(f"\n[查询结果] {len(rows)}条记录")
                        print(f"\n[前5条记录]\n{rows[:5]}")
                        time.sleep(2)

                        # 生成最终解释
                        final_response = call_ai_for_final_answer(
                            result_interpreter,
                            question,
                            "\n".join(str(row) for row in rows[:20]),  # 只展示前5条
                            sql_code
                        )
                        final_response = re.sub(r'<think>.*?</think>', '', final_response, flags=re.DOTALL).strip()
                        final_data = parse_json_response(final_response)
                        print(f"\n[最终回答]\n{final_data['message']}")

                        # 保存上下文
                        memory.save_context(
                            {"question": question},
                            {"answer": f"{final_data['message']}\nSQL: {sql_code}"}
                        )
                    else:
                        time.sleep(2)
                        # 生成最终解释
                        final_response = call_ai_for_final_answer(
                            result_interpreter,
                            question,
                            f"[提示] 已处理 {len(rows)}条记录",
                            sql_code
                        )
                        final_response = re.sub(r'<think>.*?</think>', '', final_response, flags=re.DOTALL).strip()
                        final_data = parse_json_response(final_response)
                        print(f"\n[最终回答]\n{final_data['message']}")
                        memory.save_context(
                            {"question": question},
                            {"answer": "该查询未返回有效结果"}
                        )
            else:
                print(f"\n[AI回答]\n{sql_data['message']}")
                memory.save_context(
                    {"question": question},
                    {"answer": sql_data["message"]}
                )

        except Exception as e:
            print(f"\n[错误] {str(e)}")
            memory.save_context(
                {"question": question},
                {"answer": f"处理时发生错误：{str(e)}"}
            )


if __name__ == "__main__":
    main()
